/**
 * @file main.cpp
 * @author SDK_TEAM
 * @brief
 * @version 0.1
 * @date 2022-10-26
 *
 * Copyright:
 * © 2018 北京灵汐科技有限公司 版权所有。
 * 注意：以下内容均为北京灵汐科技有限公司原创，未经本公司允许，不得转载，否则将视为侵权；对于不遵守此声明或者其他违法使用以下内容者，本公司依法保留追究权。
 * © 2018 Lynxi Technologies Co., Ltd. All rights reserved.
 * NOTICE: All information contained here is, and remains the property of Lynxi.
 * This file can not be copied or distributed without the permission of Lynxi
 * Technologies Co., Ltd.
 *
 */

#include <dlfcn.h>
#include <lyn_api.h>
#include <nlohmann/json.hpp>
#include <string>
#include "imageDecoder.h"
#include "imageEncoder.h"
#include "ipeParamYolov5s6.h"
#include "opencvWindow.hpp"
#include "osd_plugin.h"
#include "util.hpp"
#include "yolox_post_process.h"
#include "CLI/CLI11.hpp"

using namespace std;

using YOLOX_POST_FUNC = int (*)(YoloxPostProcessInfo_t *);

int main(int argc, char *argv[])
{
  CLI::App app;
  std::string inputFile;
  std::string outputFile;
  int deviceId;
  std::string postPlugin;
  int pluginType;
  std::string yolov5s6Model;
  std::string labelFile;
  std::string osdPlugin;

  app.add_option("input_image", inputFile, "input jpeg file path")->default_val("../data/yolov5s6_detect_with_image_sample/test.jpeg");
  app.add_option("output_image", outputFile, "output jpeg file path")->default_val("../output.jpeg");
  app.add_option("device_id", deviceId, "lynxi device id")->default_val(0);
  app.add_option("-p,--post_plugin", postPlugin, "yolov5s6 post process plugin path")->default_val("../common/plugin/postprocess/cross_compiled_lib/libYolov5n6X7PostPlugin.so");
  app.add_option("-t,--plugin_type", pluginType, "post process plugin compile type, 0: not cross compile, 1: cross compile")->default_val(1);
  app.add_option("-m,--yolov5s6_model", yolov5s6Model, "yolov5s6 model path")->default_val("../model/yolov5s6_b1_h1280_w1280_c3_iuint8_ofloat16_0000_x7_quant/Net_0");
  app.add_option("-l,--label_file", labelFile, "yolov5s6 label path")->default_val("../common/plugin/postprocess/yolox/label/label.v5x7.json");
  app.add_option("-o,--osd_plugin", osdPlugin, "osd plugin path")->default_val("../common/plugin/osd/lib/libOsdPlugin.so");

  // 解析命令行参数
  CLI11_PARSE(app, argc, argv);

  if (!IsFileExist(inputFile))
  {
    cout << "error : input jpeg file is not exsit." << endl;
    return 0;
  }

  if (!IsValidDevice(deviceId))
  {
    cout << "error : device id is invalid." << endl;
    return 0;
  }

  if (!IsFileExist(postPlugin))
  {
    cout << "error : yolov5s6 post process plugin is not exsit." << endl;
    return 0;
  }

  if (pluginType != 0 && pluginType != 1)
  {
    cout << "error : plugin compile type is invalid." << endl;
    return 0;
  }

  if (!IsFileExist(yolov5s6Model))
  {
    cout << "error : yolov5s6 model is not exsit." << endl;
    return 0;
  }

  if (!IsFileExist(labelFile))
  {
    cout << "error : label file is not exsit." << endl;
    return 0;
  }

  if (!IsFileExist(osdPlugin))
  {
    cout << "error : osd plugin is not exsit." << endl;
    return 0;
  }

  ifstream fJson(labelFile);
  stringstream ss;
  ss << fJson.rdbuf();
  nlohmann::json js = nlohmann::json::parse(ss.str());

  int labelCount = js.size();
  if (labelCount == 0)
  {
    cout << "error : label file is not invalid." << endl;
    return 0;
  }

  LabelList labelList;
  labelList.labelNum = labelCount;
  labelList.labels = new LabelInfo[labelCount];
  memset(labelList.labels, 0, sizeof(LabelInfo) * labelCount);
  string labelName;
  for (int i = 0; i < labelCount; ++i)
  {
    labelName = js[i]["desc"].get<string>();
    strncpy(labelList.labels[i].name, labelName.c_str(),
            sizeof(labelList.labels[i].name) - 1);
  }

  // 2. 创建 context 与 stream
  lynContext_t ctx = nullptr;
  CHECK_ERR(lynCreateContext(&ctx, deviceId));
  CHECK_ERR(lynRegisterErrorHandler(StreamErrorHandler, nullptr));
  lynStream_t stream = nullptr;
  CHECK_ERR(lynCreateStream(&stream));

  // 3. 初始化解码图片类
  ImageDecoder imageDec;
  imageDec.Init(inputFile);
  lynImageInfo_t imgInfo;
  imageDec.GetImageDecInfo(&imgInfo);

  // 4. 同步解码图片到 Device 侧
  cout << "decode image" << endl;
  lynCodecBuf_t decodeImg;
  decodeImg.size = imgInfo.output.predictBufSize;
  CHECK_ERR(lynMalloc((void **)&decodeImg.data, decodeImg.size));
  imageDec.DecodeImageToDevice(stream, decodeImg.data);

  // 5. 加载模型
  ModelInfo yolov5s6ModelInfo;
  yolov5s6ModelInfo.LoadModelByPath(yolov5s6Model.c_str());
  lynPixelFormat_t pixelFormat = LYN_PIX_FMT_NV12;

  // 6. IPE 处理
  cout << "yolov5s6 detect inference" << endl;
  uint8_t *pIpeBufOut = nullptr;
  CHECK_ERR(lynMalloc((void **)&pIpeBufOut, yolov5s6ModelInfo.inputSize));
  CHECK_ERR(lynMemset(pIpeBufOut, 0, yolov5s6ModelInfo.inputSize));
  IpeParamYolov5s6 ipeYolov5s6(yolov5s6ModelInfo.width, yolov5s6ModelInfo.height);
  ipeYolov5s6.SetImgInfo(imgInfo.output.width, imgInfo.output.height,
                         pixelFormat);
  ipeYolov5s6.CalcParam(stream, decodeImg.data, pIpeBufOut);

  // 7. 调用 APU 推理接口
  uint8_t *pDevApuOutData = nullptr;
  CHECK_ERR(lynMalloc((void **)&pDevApuOutData, yolov5s6ModelInfo.outputSize));
  CHECK_ERR(lynExecuteModelAsync(stream, yolov5s6ModelInfo.model, pIpeBufOut,
                                 pDevApuOutData, yolov5s6ModelInfo.batchSize));

  // 8. 加载post Plugin, 进行模型后处理
  lynBoxesInfo hostBoxesInfo;
  YoloxPostProcessInfo_t post_info;
  memset(&post_info, 0, sizeof(YoloxPostProcessInfo_t));

  post_info.is_pad_resize = 1;
  post_info.score_threshold = 0.25;
  post_info.nms_threshold = 0.45;
  post_info.nms_top_k = 500;
  post_info.anchorSize = 3;
  post_info.width = yolov5s6ModelInfo.width;
  post_info.height = yolov5s6ModelInfo.height;
  post_info.ori_width = imgInfo.output.width;
  post_info.ori_height = imgInfo.output.height;
  post_info.class_num = labelCount;
  lynBoxesInfo *pDevBoxesInfo;
  CHECK_ERR(lynMalloc((void **)&pDevBoxesInfo, sizeof(lynBoxesInfo)));

  if (pluginType == 0)
  { // 动态库加载
    void *handle = dlopen(postPlugin.c_str(), RTLD_NOW | RTLD_GLOBAL);
    if (handle == nullptr)
    {
      cout << "error : post plugin load fail." << endl;
      return 0;
    }
    YOLOX_POST_FUNC yoloxPostFunc =
        (YOLOX_POST_FUNC)dlsym(handle, "lynYoloxPostProcess");
    if (yoloxPostFunc == nullptr)
    {
      cout << "error : post plugin has no lynYoloxPostProcess function."
           << endl;
      dlclose(handle);
      return 0;
    }

    lynSynchronizeStream(stream);
    uint8_t *pHostApuOutData = new uint8_t[yolov5s6ModelInfo.outputSize];
    CHECK_ERR(lynMemcpy(pHostApuOutData, pDevApuOutData,
                        yolov5s6ModelInfo.outputSize, ServerToClient));

    post_info.output_tensor = pHostApuOutData;
    post_info.labelList = labelList;

    lynBoxesInfo *boxesInfo = new lynBoxesInfo;
    post_info.boxesInfo = boxesInfo;

    yoloxPostFunc(&post_info);

    CHECK_ERR(lynMemcpy(pDevBoxesInfo, boxesInfo, sizeof(lynBoxesInfo),
                        ClientToServer));

    delete boxesInfo;
    delete[] pHostApuOutData;
    dlclose(handle);
  }
  else
  { // lynxi plugin加载
    lynPlugin_t plugin;
    CHECK_ERR(lynPluginRegister(&plugin, postPlugin.c_str()));

    post_info.output_tensor = pDevApuOutData;

    LabelInfo *devLabels = nullptr;
    CHECK_ERR(
        lynMalloc((void **)&devLabels, sizeof(LabelInfo) * labelList.labelNum));
    CHECK_ERR(lynMemcpy(devLabels, labelList.labels,
                        sizeof(LabelInfo) * labelList.labelNum,
                        ClientToServer));
    post_info.labelList.labelNum = labelList.labelNum;
    post_info.labelList.labels = devLabels;

    post_info.boxesInfo = (lynBoxesInfo *)pDevBoxesInfo;

    CHECK_ERR(lynPluginRunAsync(stream, plugin, "lynYoloxPostProcess",
                                &post_info, sizeof(post_info)));
    lynSynchronizeStream(stream);

    CHECK_ERR(lynFree(devLabels));
    CHECK_ERR(lynPluginUnregister(plugin));
  }

  // 9. 加载osd Plugin, 将结果叠加到图像上
  {
    lynPlugin_t plugin;
    CHECK_ERR(lynPluginRegister(&plugin, osdPlugin.c_str()));

    lynDrawBoxAndTextPara para;
    para.imgData = decodeImg.data;
    para.imgFmt = LYN_PIX_FMT_NV12;
    para.imgW = imgInfo.output.width;
    para.imgH = imgInfo.output.height;
    para.boxesInfo = pDevBoxesInfo;
    para.boxColor = DRAW_COLOR_BLUE;
    para.boxThick = DRAW_THICK_2;
    para.fontSize = FONT_SIZE_24;
    para.fontColor = DRAW_COLOR_BLUE;
    CHECK_ERR(lynPluginRunAsync(stream, plugin, "lynDrawBoxAndText", &para,
                                sizeof(para)));
    lynSynchronizeStream(stream);
    CHECK_ERR(lynFree(pDevBoxesInfo));
    CHECK_ERR(lynPluginUnregister(plugin));
  }

  // 等待 plugin 处理完，释放内存
  CHECK_ERR(lynFree(pDevApuOutData));
  CHECK_ERR(lynFree(pIpeBufOut));

  // 9. 进行图片编码
  cout << "save new image" << endl;
  ImageEncoder encoder;
  encoder.EncodeImageAndSave(stream, outputFile, decodeImg, imgInfo);

  // 10. 销毁资源
  CHECK_ERR(lynDestroyStream(stream));
  CHECK_ERR(lynFree(decodeImg.data));
  yolov5s6ModelInfo.UnLoadModel();

  if (ctx)
  {
    CHECK_ERR(lynDestroyContext(ctx));
  }

  delete[] labelList.labels;

  return 0;
}
